import json
import sys
import pymysql

from graphviz import Digraph


def process(num):
    """
    :param num:命令行输入参数为家谱id
    :return:返回查询的json文件
    """
    try:
        conn = pymysql.Connect(host='114.55.254.213', user='guoqi',
                               passwd='123', db='family', port=3306, charset='utf8')
        cur = conn.cursor()
        sql = "SELECT person_id,father_id,last_name,first_name,child_type,ranking," \
              "name from family_person WHERE family_archive_id=" + \
              str(num) + ";"
        cur.execute(sql)
        data = cur.fetchall()
        cur.close()
        conn.close()
        jsonData = []
        # 循环读取元组数据
        for row in data:
            result = {'person_id': str(row[0]), 'father_id': str(row[1]), 'last_name': str(row[2]),
                      'first_name': str(row[3]), 'child_type': str(row[4]), 'ranking': str(row[5]), 'name': str(row[6])}
            jsonData.append(result)
    except:
        pass
    else:
        jsondata = json.dumps(jsonData, ensure_ascii=False)
        return jsondata


def check(str1, str2):
    """
    :param str1: 孩子类型
    :param str2: 孩子排行
    :return: 简单处理后的关系
    """
    if str1 == "亲生":
        if str2 != "":
            return "第" + str2 + "子"
    else:
        return "继子"


def handle(json_data):
    """
    :param json_data:原始json文件
    :return:经过处理后的json文件
    """
    new_json = {}
    json_data = json.loads(json_data)
    for data in json_data:
        new_json[str(data['person_id'])] = {
            "father_id": data['father_id'],
            "last_name": data['last_name'],
            "first_name": data['first_name'],
            "relation": check(data['child_type'], data['ranking']),
            "name": data['first_name'] + data['last_name']
        }
    return new_json


def tree_image(num):
    """
    :param num: 命令行参数
    :return: 无返回
    """
    json_data1 = process(num)
    json_data = handle(json_data1)
    f = Digraph('familyTree', filename='family.gv', format="png")
    f.attr(rankdir='TB', size='20,5', splines='ortho')
    f.attr('node', shape='box', fontname="Sans Not-Rotated 14")
    # root_data = json_data[str(1)]['name']
    # if(root_data=="None"):
    #     root_data="ss"

    usable = []
    for data in json_data:
        usable.append(data)
    # f.node(root_data)
    # print(usable)
    for i in usable:
        if json_data[i]["father_id"] == '0':
            root_data = json_data[str(json_data[i])]['name']
            f.node(root_data)
            usable.remove(i)
    try:
        for i in usable:
            f.edge(json_data[str(json_data[i]["father_id"])]["name"], json_data[i]["name"],
                   label=json_data[i]["relation"], fontname="Sans Not-Rotated 8")
        # f.view()
        f.save()
    except KeyError:
        pass


if __name__ == '__main__':
    tree_image(sys.argv[1])
